{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ed7178e-1a9c-4784-9ae8-3dd125db0146",
   "metadata": {},
   "source": [
    "# Mistral Ai Chat - MistralClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "945cd540-b243-407e-b5bc-347839ea75c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: mistralai\n",
      "Version: 0.1.6\n",
      "Summary: \n",
      "Home-page: \n",
      "Author: Bam4d\n",
      "Author-email: bam4d@mistral.ai\n",
      "License: \n",
      "Location: /opt/conda/lib/python3.11/site-packages\n",
      "Requires: httpx, orjson, pandas, pyarrow, pydantic\n",
      "Required-by: \n"
     ]
    }
   ],
   "source": [
    "!pip show mistralai"
   ]
  },
{
   "cell_type": "markdown",
   "id": "53b0140b-97b6-4339-8a1b-a6740ad4107d",
   "metadata": {},
   "source": [
    "## Create MistralClient "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a9cfc7ca-4d18-47d9-9803-cfedb09e3f09",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from mistralai.client import MistralClient\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = MistralClient(api_key=os.environ[\"MISTRAL_API_KEY\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d161b155-1832-4800-aa90-2001b59b2cfa",
   "metadata": {},
   "source": [
    "## Load instruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "534ff961-7a09-4287-954d-458d8209e6af",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 1048.05it/s]\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import DirectoryLoader, TextLoader\n",
    "\n",
    "loader = DirectoryLoader('structurizr/llm', glob=\"instruction.txt\", show_progress=True, use_multithreading=True, loader_cls=TextLoader)\n",
    "\n",
    "instruction = loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "00f7037d-81fc-4710-920a-7387d07ae1b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are an expert on structurizr which is a domain specific language (DSL) for software architecture modeling and documenation.\n",
      "You will be provided a keyword in structurizr DSL, you need to provide the keyword gramma precisely and consistently.\n",
      "You need to check the permited children of the keyword.\n",
      "You must strictly adhere to the keyword gramma to generate the output.\n",
      "You need to embed the permited children of the keyword in the gramma.\n",
      "You do not need to provide the permited children out of the gramma.\n",
      "You need to provide the gramma, you do not need to provide any explanations.\n",
      "You do not need to privde code examples or explanations. \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(instruction[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c427e7c-9334-4d33-926c-979cf41cc306",
   "metadata": {},
   "source": [
    "## Retrieval\n",
    "\n",
    "please refer to ollama-embeding.ipynb for embeding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6d751631-d27e-4c50-b05c-dea0b637d4d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 287 ms, sys: 37.8 ms, total: 325 ms\n",
      "Wall time: 723 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "import utility\n",
    "embeddings = utility.get_embeddings()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7234f9-2c1f-4fba-a95d-22839abfe82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import pyarrow as pa\n",
    "import lancedb\n",
    "from langchain_community.vectorstores import LanceDB\n",
    "\n",
    "db = lancedb.connect(\"lancedb\")\n",
    "\n",
    "language = db.open_table('language')\n",
    "\n",
    "language_vectorstore = LanceDB(language, embeddings)\n",
    "\n",
    "retriever = language_vectorstore.as_retriever(search_kwargs={\"k\": 1})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c8e292-7b1e-48db-ade9-ae7becec7750",
   "metadata": {},
   "source": [
    "## ChatMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6ac3bed8-57ab-4b07-b2fe-78ab912f72c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "keyword = \"workspace\"\n",
    "\n",
    "docs = retriever.get_relevant_documents(keyword)\n",
    "\n",
    "language_gramma = f\"The gramma of {keyword} is provided here: {docs[0].page_content}\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4184632a-eee3-4ce4-8f86-c9071e3f0dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mistralai.models.chat_completion import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=instruction[0].page_content),\n",
    "    ChatMessage(role=\"system\", content=language_gramma),\n",
    "    ChatMessage(role=\"user\", content=keyword)\n",
    "]"
   ]
  },
  
  {
   "cell_type": "markdown",
   "id": "0d9b7df9-ac35-4107-a246-365c0428ebdc",
   "metadata": {},
   "source": [
    "## Chat with \"open-mistral-7b\" using instruction and grammars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "37dba09c-cb6f-4ab5-b441-9ad6198629ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```\n",
      "workspace [name] [description] {\n",
      "  (name: <name>)?\n",
      "  (description: <description>)?\n",
      "  (properties)?\n",
      "  (!docs)?\n",
      "  (!adrs)?\n",
      "  (!identifiers)?\n",
      "  (!impliedRelationships)?\n",
      "  model {\n",
      "    ...\n",
      "  }\n",
      "  views {\n",
      "    ...\n",
      "  }\n",
      "  (configuration)?\n",
      "}\n",
      "\n",
      "workspace extends <file|url> {\n",
      "  (name: <name>)?\n",
      "  (description: <description>)?\n",
      "  (properties)?\n",
      "  (!docs)?\n",
      "  (!adrs)?\n",
      "  (!identifiers)?\n",
      "  (!impliedRelationships)?\n",
      "  model {\n",
      "    ...\n",
      "  }\n",
      "  views {\n",
      "    ...\n",
      "  }\n",
      "  (configuration)?\n",
      "}\n",
      "```\n",
      "prompt_tokens=455 total_tokens=640 completion_tokens=185\n",
      "CPU times: user 23.6 ms, sys: 7.43 ms, total: 31 ms\n",
      "Wall time: 4.18 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = client.chat(\n",
    "    model=\"open-mistral-7b\",\n",
    "    messages=messages,\n",
    ")\n",
    "\n",
    "print(chat_response.choices[0].message.content, flush=True)\n",
    "\n",
    "print(chat_response.usage, flush=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8363175-b0d7-4ee0-be4e-0300fd36492c",
   "metadata": {},
   "source": [
    "## Chat with \"mistral-large-2402\" using instruction and grammars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "31819ff8-7957-4bc7-a5e3-4a8be8423cec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The gramma for the `workspace` keyword in structurizr DSL is as follows:\n",
      "\n",
      "```\n",
      "workspace [name] [description] {\n",
      "  (name <name>)？\n",
      "  (description <description>)？\n",
      "  (properties)？\n",
      "  (!docs)？\n",
      "  (!adrs)？\n",
      "  (!identifiers)？\n",
      "  (!impliedRelationships)？\n",
      "  (model)？\n",
      "  (views)？\n",
      "  (configuration)？\n",
      "}\n",
      "```\n",
      "\n",
      "or\n",
      "\n",
      "```\n",
      "workspace extends <file|url> {\n",
      "  (name <name>)？\n",
      "  (description <description>)？\n",
      "  (properties)？\n",
      "  (!docs)？\n",
      "  (!adrs)？\n",
      "  (!identifiers)？\n",
      "  (!impliedRelationships)？\n",
      "  (model)？\n",
      "  (views)？\n",
      "  (configuration)？\n",
      "}\n",
      "```\n",
      "\n",
      "The `<name>`, `<description>`, `<file|url>` are placeholders and should be replaced with actual values. The `name`, `description`, `properties`, `!docs`, `!adrs`, `!identifiers`, `!impliedRelationships`, `model`, `views`, and `configuration` are permitted children of the `workspace` keyword. The question mark `?` indicates that the preceding element is optional. The exclamation mark `!` before a keyword indicates that it is a block keyword and its content should be enclosed in curly braces `{}`.\n",
      "prompt_tokens=450 total_tokens=784 completion_tokens=334\n",
      "CPU times: user 31.4 ms, sys: 5.49 ms, total: 36.9 ms\n",
      "Wall time: 8.39 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = client.chat(\n",
    "    model=\"mistral-large-2402\",\n",
    "    messages=messages,\n",
    ")\n",
    "\n",
    "print(chat_response.choices[0].message.content, flush=True)\n",
    "\n",
    "print(chat_response.usage, flush=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89f5a67-ff55-4f52-a335-f1f787377ddc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
